import os
import time
import numpy as np
import cv2
import tkinter as tk
from tkinter import filedialog, Scale, Button, Label, Frame, IntVar, StringVar, DoubleVar, Canvas, Checkbutton
from PIL import Image, ImageTk
import matplotlib.pyplot as plt
from paddleocr import PaddleOCR
from skimage import measure
import torch
import torchvision
from segment_anything import sam_model_registry, SamPredictor

class SAMMaskGenerator:
    """使用PaddleOCR和SAM模型进行艺术字提取的核心处理类"""
    
    def __init__(self, use_gpu=False, model_type="vit_b"):
        self.use_gpu = use_gpu and torch.cuda.is_available()
        self.device = "cuda" if self.use_gpu else "cpu"
        
        # 初始化OCR模型
        self.ocr = PaddleOCR(
            use_angle_cls=True,       # 使用方向分类器
            lang="ch",                # 中文模型
            use_gpu=self.use_gpu,     # 是否使用GPU
            det_model_dir=None,       # 使用默认检测模型
            rec_model_dir=None,       # 使用默认识别模型
            show_log=False            # 不显示日志
        )
        
        # 初始化SAM模型
        self.model_type = model_type
        self.sam = None
        self.predictor = None
        self.sam_available = False
        try:
            self._initialize_sam()
            self.sam_available = True
        except Exception as e:
            print(f"SAM模型初始化失败，将使用备用方法: {str(e)}")
            print("请尝试安装Microsoft Visual C++ Redistributable: https://aka.ms/vs/16/release/vc_redist.x64.exe")
            
        # 图像相关属性
        self.image = None
        self.original = None
        self.mask = None
        self.height = 0
        self.width = 0
        self.result = None
        self.boxes = []
        self.prompt_points = []
        self.sam_masks = []
        
    def _initialize_sam(self):
        """初始化SAM模型"""
        try:
            # 模型类型和下载地址映射
            model_urls = {
                "vit_h": "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth",
                "vit_l": "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_l_0b3195.pth",
                "vit_b": "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth"
            }
            
            # 下载并加载模型（首次运行时会自动下载）
            print(f"正在加载SAM {self.model_type}模型，首次运行可能需要下载模型文件...")
            self.sam = sam_model_registry[self.model_type](checkpoint=model_urls[self.model_type])
            self.sam.to(device=self.device)
            self.predictor = SamPredictor(self.sam)
            print("SAM模型加载完成！")
        except Exception as e:
            print(f"加载SAM模型失败: {str(e)}")
            raise
        
    def load_image(self, image_path):
        """加载图像并保存原始副本"""
        try:
            # 使用PIL加载图像，然后转换为OpenCV格式
            pil_image = Image.open(image_path)
            self.image = cv2.cvtColor(np.array(pil_image), cv2.COLOR_RGB2BGR)
            self.original = self.image.copy()
            self.height, self.width = self.image.shape[:2]
            self.result = None
            self.boxes = []
            self.mask = None
            self.prompt_points = []
            self.sam_masks = []
            
            # 为SAM预处理图像
            if self.predictor is not None:
                self.predictor.set_image(self.image)
                
            return self.image
        except Exception as e:
            raise ValueError(f"无法读取图像: {image_path}，错误: {str(e)}")
        
    def detect_text(self):
        """使用PaddleOCR检测文本区域"""
        if self.original is None:
            raise ValueError("请先加载图像")
            
        # 运行PaddleOCR检测
        start_time = time.time()
        self.result = self.ocr.ocr(self.original, cls=True)
        process_time = time.time() - start_time
        
        # 提取文本框坐标
        self.boxes = []
        if self.result and len(self.result) > 0:
            for line in self.result:
                if isinstance(line, list) and line:
                    for item in line:
                        if isinstance(item, list) and len(item) >= 1:
                            if isinstance(item[0], list) and len(item[0]) == 4:
                                # 文本框坐标 [[x1,y1], [x2,y2], [x3,y3], [x4,y4]]
                                self.boxes.append(np.array(item[0], dtype=np.int32))
        
        return self.boxes, process_time
    
    def _generate_prompt_points(self):
        """从检测到的文本框生成SAM的提示点"""
        if not self.boxes:
            return []
            
        prompt_points = []
        labels = []  # 1表示前景点，0表示背景点
        
        for box in self.boxes:
            # 计算多边形框的中心点作为前景点
            center_x = np.mean(box[:, 0])
            center_y = np.mean(box[:, 1])
            prompt_points.append([center_x, center_y])
            labels.append(1)  # 标记为前景点
            
            # 在框的外部添加背景点
            padding = 10  # 背景点与框的距离
            
            # 获取框的最小最大坐标
            x_min, y_min = np.min(box, axis=0)
            x_max, y_max = np.max(box, axis=0)
            
            # 在框的四周添加背景点
            bg_points = [
                [x_min - padding, center_y],  # 左边
                [x_max + padding, center_y],  # 右边
                [center_x, y_min - padding],  # 上边
                [center_x, y_max + padding]   # 下边
            ]
            
            # 过滤掉越界的点
            for bg_point in bg_points:
                x, y = bg_point
                if 0 <= x < self.width and 0 <= y < self.height:
                    prompt_points.append(bg_point)
                    labels.append(0)  # 标记为背景点
        
        self.prompt_points = np.array(prompt_points)
        return np.array(prompt_points), np.array(labels)
        
    def create_sam_mask(self, confidence_threshold=0.0, stability_threshold=0.0, 
                        overlap_threshold=0.5, min_area=100):
        """使用SAM模型根据提示点生成分割掩码，如果SAM不可用则使用备用方法"""
        if self.original is None or not self.boxes:
            raise ValueError("请先加载图像并运行文本检测")
        
        # 如果SAM模型不可用，使用备用方法
        if not self.sam_available or self.predictor is None:
            return self._create_fallback_mask(min_area)
            
        # SAM可用，使用标准流程
        # 生成提示点和标签
        points, labels = self._generate_prompt_points()
        
        if len(points) == 0:
            raise ValueError("未生成有效的提示点")
            
        # 使用SAM模型生成掩码
        masks, scores, _ = self.predictor.predict(
            point_coords=points,
            point_labels=labels,
            multimask_output=True,  # 生成多个候选掩码
            mask_input=None,        # 可选的初始掩码
        )
        
        # 选择置信度最高的掩码
        best_mask_idx = np.argmax(scores)
        best_mask = masks[best_mask_idx]
        
        # 转换为OpenCV格式的掩码
        mask = best_mask.astype(np.uint8) * 255
        
        # 过滤小区域
        if min_area > 0:
            # 标记连通区域
            labels = measure.label(mask)
            properties = measure.regionprops(labels)
            
            # 过滤小区域
            for prop in properties:
                if prop.area < min_area:
                    mask[labels == prop.label] = 0
        
        self.mask = mask
        return mask
    
    def _create_fallback_mask(self, min_area=100):
        """当SAM不可用时的备用掩码生成方法"""
        print("使用备用方法生成蒙版")
        # 创建空白蒙版
        mask = np.zeros((self.height, self.width), dtype=np.uint8)
        
        # 对每个检测到的文本框区域进行处理
        for box in self.boxes:
            # 计算文本框的边界
            x_min, y_min = np.min(box, axis=0)
            x_max, y_max = np.max(box, axis=0)
            
            # 确保边界在图像范围内
            x_min = max(0, x_min)
            y_min = max(0, y_min)
            x_max = min(self.width - 1, x_max)
            y_max = min(self.height - 1, y_max)
            
            # 提取ROI区域
            roi = self.original[y_min:y_max+1, x_min:x_max+1]
            if roi.size == 0:  # 跳过空区域
                continue
                
            # 转换为灰度
            roi_gray = cv2.cvtColor(roi, cv2.COLOR_BGR2GRAY) if len(roi.shape) == 3 else roi.copy()
                
            # 应用自适应阈值提取文字
            block_size = 11
            c = 2
            roi_thresh = cv2.adaptiveThreshold(
                roi_gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, 
                cv2.THRESH_BINARY_INV, block_size, c
            )
                
            # 过滤小区域
            if min_area > 0:
                # 找到所有轮廓
                contours, _ = cv2.findContours(
                    roi_thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
                )
                # 在临时蒙版上填充足够大的轮廓
                temp_mask = np.zeros_like(roi_thresh)
                for contour in contours:
                    if cv2.contourArea(contour) > min_area:
                        cv2.drawContours(temp_mask, [contour], -1, 255, -1)
                roi_thresh = temp_mask
                
            # 将处理后的ROI放回原始蒙版
            mask[y_min:y_max+1, x_min:x_max+1] = cv2.bitwise_or(
                mask[y_min:y_max+1, x_min:x_max+1], roi_thresh
            )
        
        self.mask = mask
        return mask
    
    def visualize_detection(self):
        """可视化检测结果"""
        if self.original is None or not self.boxes:
            raise ValueError("请先加载图像并运行文本检测")
            
        # 克隆原始图像以在上面绘制
        visualization = self.original.copy()
        
        # 绘制检测到的文本框
        for box in self.boxes:
            cv2.polylines(visualization, [box], True, (0, 255, 0), 2)
            
        return visualization
    
    def visualize_prompt_points(self):
        """可视化提示点"""
        if self.original is None or len(self.prompt_points) == 0:
            raise ValueError("请先加载图像并生成提示点")
            
        # 克隆原始图像以在上面绘制
        visualization = self.original.copy()
        
        # 绘制提示点
        for i, point in enumerate(self.prompt_points):
            x, y = point
            # 确保坐标是整数
            x, y = int(x), int(y)
            # 确保坐标在图像范围内
            if 0 <= x < self.width and 0 <= y < self.height:
                # 前景点绘制为绿色，背景点绘制为红色
                color = (0, 255, 0) if i < len(self.boxes) else (0, 0, 255)
                # 绘制点
                cv2.circle(visualization, (x, y), 5, color, -1)  # -1表示实心圆
                
        return visualization
    
    def save_mask(self, output_path):
        """保存蒙版图像"""
        if self.mask is None:
            raise ValueError("请先创建蒙版")
        
        cv2.imwrite(output_path, self.mask)
        return output_path


class MaskGeneratorUI:
    """用户界面类，提供交互调整参数的界面"""
    
    def __init__(self, root):
        self.root = root
        self.root.title("艺术字蒙版生成器")
        self.root.geometry("1200x800")
        
        # 根据GPU可用性决定是否使用GPU
        use_gpu = False
        try:
            use_gpu = torch.cuda.is_available()
        except:
            pass
            
        self.processor = SAMMaskGenerator(use_gpu=use_gpu)
        self.image_path = ""
        self.setup_ui()
        
    def setup_ui(self):
        """设置UI组件"""
        # 创建左右布局框架
        left_frame = Frame(self.root)
        left_frame.pack(side=tk.LEFT, fill=tk.BOTH, expand=True)
        
        right_frame = Frame(self.root)
        right_frame.pack(side=tk.RIGHT, fill=tk.BOTH)
        
        # 图像显示区域
        self.image_label = Label(left_frame, text="原始图像")
        self.image_label.pack()
        
        self.canvas_original = Label(left_frame)
        self.canvas_original.pack(padx=10, pady=5)
        
        self.detect_label = Label(left_frame, text="检测结果")
        self.detect_label.pack()
        
        self.canvas_detection = Label(left_frame)
        self.canvas_detection.pack(padx=10, pady=5)
        
        self.prompt_label = Label(left_frame, text="提示点")
        self.prompt_label.pack()
        
        self.canvas_prompt = Label(left_frame)
        self.canvas_prompt.pack(padx=10, pady=5)
        
        self.mask_label = Label(left_frame, text="蒙版结果")
        self.mask_label.pack()
        
        self.canvas_mask = Label(left_frame)
        self.canvas_mask.pack(padx=10, pady=5)
        
        # 控制面板
        control_frame = Frame(right_frame)
        control_frame.pack(fill=tk.BOTH, padx=10, pady=10)
        
        # 文件操作按钮
        file_frame = Frame(control_frame)
        file_frame.pack(fill=tk.X, pady=5)
        
        self.load_btn = Button(file_frame, text="加载图像", command=self.load_image)
        self.load_btn.pack(side=tk.LEFT, padx=5)
        
        self.detect_btn = Button(file_frame, text="检测文本", command=self.detect_text)
        self.detect_btn.pack(side=tk.LEFT, padx=5)
        
        # 根据SAM可用性决定是否显示提示点按钮
        if self.processor.sam_available:
            self.prompt_btn = Button(file_frame, text="生成提示点", command=self.generate_prompts)
            self.prompt_btn.pack(side=tk.LEFT, padx=5)
        
        self.mask_btn = Button(file_frame, text="生成蒙版", command=self.create_mask)
        self.mask_btn.pack(side=tk.LEFT, padx=5)
        
        self.save_btn = Button(file_frame, text="保存蒙版", command=self.save_mask)
        self.save_btn.pack(side=tk.LEFT, padx=5)
        
        # 参数调整区域
        params_frame = Frame(control_frame)
        params_frame.pack(fill=tk.BOTH, expand=True, pady=10)
        
        # SAM参数设置
        Label(params_frame, text="SAM模型参数", font=("Arial", 10, "bold")).pack(anchor="w")
        
        # 置信度阈值
        conf_frame = Frame(params_frame)
        conf_frame.pack(fill=tk.X, pady=2)
        Label(conf_frame, text="置信度阈值:").pack(side=tk.LEFT)
        self.conf_var = DoubleVar(value=0.0)
        Scale(conf_frame, from_=0.0, to=1.0, resolution=0.01, orient=tk.HORIZONTAL, 
              variable=self.conf_var).pack(side=tk.LEFT, fill=tk.X, expand=True)
        
        # 稳定性阈值
        stab_frame = Frame(params_frame)
        stab_frame.pack(fill=tk.X, pady=2)
        Label(stab_frame, text="稳定性阈值:").pack(side=tk.LEFT)
        self.stab_var = DoubleVar(value=0.0)
        Scale(stab_frame, from_=0.0, to=1.0, resolution=0.01, orient=tk.HORIZONTAL, 
              variable=self.stab_var).pack(side=tk.LEFT, fill=tk.X, expand=True)
        
        # 重叠阈值
        overlap_frame = Frame(params_frame)
        overlap_frame.pack(fill=tk.X, pady=2)
        Label(overlap_frame, text="重叠阈值:").pack(side=tk.LEFT)
        self.overlap_var = DoubleVar(value=0.5)
        Scale(overlap_frame, from_=0.0, to=1.0, resolution=0.01, orient=tk.HORIZONTAL, 
              variable=self.overlap_var).pack(side=tk.LEFT, fill=tk.X, expand=True)
        
        # 最小区域面积
        area_frame = Frame(params_frame)
        area_frame.pack(fill=tk.X, pady=2)
        Label(area_frame, text="最小区域:").pack(side=tk.LEFT)
        self.area_var = IntVar(value=100)
        Scale(area_frame, from_=0, to=1000, orient=tk.HORIZONTAL, 
              variable=self.area_var).pack(side=tk.LEFT, fill=tk.X, expand=True)
        
        # 状态信息
        self.status_var = StringVar(value="请加载图像")
        status_bar = Label(self.root, textvariable=self.status_var, bd=1, relief=tk.SUNKEN, anchor=tk.W)
        status_bar.pack(side=tk.BOTTOM, fill=tk.X)
    
    def load_image(self):
        """加载图像文件"""
        file_path = filedialog.askopenfilename(
            filetypes=[("图像文件", "*.jpg *.jpeg *.png *.bmp")],
            initialdir=os.getcwd()  # 使用当前目录作为起始目录
        )
        
        if file_path:
            try:
                self.image_path = file_path
                self.status_var.set("正在加载图像...")
                self.root.update()
                
                # 加载图像到处理器
                self.processor.load_image(file_path)
                
                # 显示图像
                self.display_image(self.processor.image, self.canvas_original)
                self.status_var.set(f"已加载图像: {os.path.basename(file_path)}")
                
                # 清除之前的结果
                self.canvas_detection.config(image="")
                self.canvas_prompt.config(image="")
                self.canvas_mask.config(image="")
                
            except Exception as e:
                self.status_var.set(f"错误: {str(e)}")
    
    def detect_text(self):
        """检测图像中的文本区域"""
        if not hasattr(self.processor, 'original') or self.processor.original is None:
            self.status_var.set("请先加载图像")
            return
            
        try:
            self.status_var.set("正在检测文本...")
            self.root.update()
            
            # 运行文本检测
            boxes, process_time = self.processor.detect_text()
            
            if not boxes:
                self.status_var.set("未检测到文本!")
                return
            
            # 显示检测结果
            visualization = self.processor.visualize_detection()
            self.display_image(visualization, self.canvas_detection)
            
            self.status_var.set(f"文本检测完成 (处理时间: {process_time:.2f}秒) - 检测到 {len(boxes)} 个区域")
        except Exception as e:
            self.status_var.set(f"检测错误: {str(e)}")
    
    def generate_prompts(self):
        """生成SAM提示点"""
        if not hasattr(self.processor, 'boxes') or not self.processor.boxes:
            self.status_var.set("请先运行文本检测")
            return
            
        try:
            self.status_var.set("生成提示点中...")
            self.root.update()
            
            # 生成提示点
            self.processor._generate_prompt_points()
            
            # 显示提示点
            visualization = self.processor.visualize_prompt_points()
            self.display_image(visualization, self.canvas_prompt)
            
            self.status_var.set(f"已生成 {len(self.processor.prompt_points)} 个提示点")
        except Exception as e:
            self.status_var.set(f"生成提示点错误: {str(e)}")
    
    def create_mask(self):
        """创建文本蒙版"""
        if not hasattr(self.processor, 'boxes') or not self.processor.boxes:
            self.status_var.set("请先运行文本检测")
            return
            
        try:
            self.status_var.set("生成蒙版中...")
            self.root.update()
            
            # 如果SAM可用且已生成提示点，则使用提示点
            if (self.processor.sam_available and 
                hasattr(self.processor, 'prompt_points') and 
                len(self.processor.prompt_points) > 0):
                
                # 获取当前参数
                params = {
                    'confidence_threshold': self.conf_var.get(),
                    'stability_threshold': self.stab_var.get(),
                    'overlap_threshold': self.overlap_var.get(),
                    'min_area': self.area_var.get()
                }
                
                mask = self.processor.create_sam_mask(**params)
            else:
                # 使用备用方法
                mask = self.processor._create_fallback_mask(self.area_var.get())
            
            # 显示蒙版
            self.display_image(mask, self.canvas_mask)
            
            self.status_var.set(f"蒙版生成完成")
        except Exception as e:
            self.status_var.set(f"蒙版生成错误: {str(e)}")
    
    def save_mask(self):
        """保存蒙版图像"""
        if not hasattr(self.processor, 'mask') or self.processor.mask is None:
            self.status_var.set("请先生成蒙版")
            return
            
        # 构建默认文件名
        if self.image_path:
            directory = os.path.dirname(self.image_path)
            basename = os.path.splitext(os.path.basename(self.image_path))[0]
            default_filename = f"{basename}_sam_mask.jpg"
            default_path = os.path.join(directory, default_filename)
        else:
            default_path = "sam_mask.jpg"
            
        file_path = filedialog.asksaveasfilename(
            defaultextension=".jpg",
            filetypes=[("JPEG 文件", "*.jpg"), ("PNG 文件", "*.png"), ("所有文件", "*.*")],
            initialfile=default_path
        )
        
        if file_path:
            try:
                cv2.imwrite(file_path, self.processor.mask)
                self.status_var.set(f"蒙版已保存至: {os.path.basename(file_path)}")
            except Exception as e:
                self.status_var.set(f"保存错误: {str(e)}")
    
    def display_image(self, cv_image, canvas):
        """在画布上显示OpenCV图像"""
        # 确保图像是灰度或BGR格式
        if len(cv_image.shape) == 2:
            # 灰度图像转RGB
            display_image = cv2.cvtColor(cv_image, cv2.COLOR_GRAY2RGB)
        else:
            # BGR转RGB
            display_image = cv2.cvtColor(cv_image, cv2.COLOR_BGR2RGB)
            
        # 调整图像大小以适应界面
        max_height = 175
        height, width = display_image.shape[:2]
        
        if height > max_height:
            scale = max_height / height
            new_width = int(width * scale)
            display_image = cv2.resize(display_image, (new_width, max_height))
            
        # 转换为Tkinter可显示的格式
        pil_image = Image.fromarray(display_image)
        tk_image = ImageTk.PhotoImage(image=pil_image)
        
        # 更新画布
        canvas.config(image=tk_image)
        canvas.image = tk_image  # 保持引用以防止被垃圾回收


def main():
    """主函数"""
    root = tk.Tk()
    app = MaskGeneratorUI(root)
    root.mainloop()


if __name__ == "__main__":
    main()
